# ==============================================================================
# 文件名: 字帖图像处理工具.py
# 功能: 基于OpenCV和Tkinter的字帖图像预处理工具，支持灰度转换、二值化、形态学操作、
#       轮廓检测及汉字提取等功能，适用于书法字帖图像的数字化处理与单字分割。
# 核心模块:
#   - 图像处理算法: 包含灰度转换、Otsu阈值、腐蚀/膨胀/闭运算、Canny边缘检测等基础操作
#   - GUI界面: 提供图像载入/重置、交互式裁剪、实时参数预览、汉字提取保存等可视化功能
# 作者: 开发者自定义
# 版本: 1.0
# ==============================================================================
import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import cv2
import numpy as np
import os

# ========== 图像处理算法 ==========
def to_gray(image):
    if len(image.shape) == 2 or (len(image.shape) == 3 and image.shape[2] == 1):
        return image
    return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

def otsu_threshold(gray_img, maxval=255):
    _, result = cv2.threshold(gray_img, 0, maxval, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return result

def erode(gray_img, ksize=3, iterations=1):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (ksize, ksize))
    return cv2.erode(gray_img, kernel, iterations=iterations)

def dilate(gray_img, ksize=3, iterations=1):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (ksize, ksize))
    return cv2.dilate(gray_img, kernel, iterations=iterations)

def morphology_close(gray_img, ksize=3, iterations=1):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (ksize, ksize))
    return cv2.morphologyEx(gray_img, cv2.MORPH_CLOSE, kernel, iterations=iterations)

def canny_edge(gray_img, th1=50, th2=150):
    return cv2.Canny(gray_img, threshold1=th1, threshold2=th2)

def merge_boxes(boxes, threshold=10):
    merged = []
    for box in boxes:
        x, y, w, h = box
        found = False
        for i, (mx, my, mw, mh) in enumerate(merged):
            if not (x > mx+mw+threshold or mx > x+w+threshold or y > my+mh+threshold or my > y+h+threshold):
                nx = min(x, mx)
                ny = min(y, my)
                nw = max(x+w, mx+mw) - nx
                nh = max(y+h, my+mh) - ny
                merged[i] = (nx, ny, nw, nh)
                found = True
                break
        if not found:
            merged.append(box)
    return merged

def find_contours_with_merge(gray_img, area_thresh=500, wh_ratio=(0.3, 3), min_wh=15, merge_thresh=10):
    contours, _ = cv2.findContours(gray_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    boxes = []
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        area = cv2.contourArea(cnt)
        if area > area_thresh and wh_ratio[0] < w / h < wh_ratio[1] and w > min_wh and h > min_wh:
            boxes.append((x, y, w, h))
    merged_boxes = merge_boxes(boxes, threshold=merge_thresh)
    if len(gray_img.shape) == 2:
        contour_img = cv2.cvtColor(gray_img, cv2.COLOR_GRAY2BGR)
    else:
        contour_img = gray_img.copy()
    for x, y, w, h in merged_boxes:
        cv2.rectangle(contour_img, (x, y), (x + w, y + h), (0, 255, 0), 2)
    return contour_img, merged_boxes

def median_filter(gray_img, ksize=3):
    return cv2.medianBlur(gray_img, ksize)

# ========== 主界面与功能 ==========
class ImageProcessingApp:
    """字帖图像处理工具主界面类
    
    功能说明:
    1. 文件操作: 支持JPG/PNG图像载入，可重置为原始图像
    2. 图像显示: 自动缩放适应窗口，支持居中显示
    3. 交互裁剪: 通过鼠标绘制矩形区域实现图像裁剪
    4. 图像处理: 包含9种常用操作（转灰度/阈值处理/腐蚀/膨胀/闭运算/Canny检测等）
    5. 高级功能: 实时参数预览（腐蚀/膨胀/闭运算时调整参数可实时查看效果）、
               汉字轮廓检测（基于Canny边缘和轮廓合并算法）、单字提取保存
    """
    def __init__(self, root):
        self.root = root
        self.root.title("字帖图像处理工具")
        self.root.geometry("800x600")
        self.current_image = None
        self.original_image = None
        self.crop_region = None
        self.is_drawing = False
        self.rect_start_x = None
        self.rect_start_y = None
        self.rect_end_x = None
        self.rect_end_y = None
        self.current_rect = None
        self.menu_bar = tk.Menu(root)
        self.file_menu = tk.Menu(self.menu_bar, tearoff=0)
        self.file_menu.add_command(label="载入图像", command=self.load_image)
        self.file_menu.add_command(label="重置为原图", command=self.reset_to_original)
        self.process_menu = tk.Menu(self.menu_bar, tearoff=0)
        self.process_menu.add_command(label="裁剪", command=self.crop_image)
        self.process_menu.add_command(label="原图转灰度图", command=lambda: self.process_image(1))
        self.process_menu.add_command(label="全局阈值处理", command=lambda: self.process_image(2))
        self.process_menu.add_command(label="应用腐蚀操作", command=lambda: self.process_image(3))
        self.process_menu.add_command(label="应用膨胀操作", command=lambda: self.process_image(4))
        self.process_menu.add_command(label="应用闭运算", command=lambda: self.process_image(6))
        self.process_menu.add_command(label="Canny边缘检测", command=lambda: self.process_image(7))
        self.process_menu.add_command(label="轮廓检测", command=self.contour_detect_with_canny)
        self.process_menu.add_command(label="提取汉字", command=self.extract_and_save_chars)
        self.process_menu.add_command(label="中值滤波去除小白点", command=self.apply_median_filter)
        self.menu_bar.add_cascade(label="文件", menu=self.file_menu)
        self.menu_bar.add_cascade(label="处理", menu=self.process_menu)
        root.config(menu=self.menu_bar)
        self.canvas = tk.Canvas(root)
        self.canvas.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        self.canvas.bind("<Button-1>", self.on_mouse_down)
        self.canvas.bind("<B1-Motion>", self.on_mouse_move)
        self.canvas.bind("<ButtonRelease-1>", self.on_mouse_up)
        self.tk_img = None

    def load_image(self):
        file_path = filedialog.askopenfilename(
            initialdir="图片存放",
            title="选择图像文件",
            filetypes=(("JPG文件", "*.jpg"), ("PNG文件", "*.png"), ("所有文件", "*.*"))
        )
        if not file_path:
            return
        try:
            with open(file_path, "rb") as f:
                img_bytes = f.read()
            img_np = np.frombuffer(img_bytes, np.uint8)
            self.original_image = cv2.imdecode(img_np, cv2.IMREAD_COLOR)
            if self.original_image is None:
                raise Exception("无法识别的图像格式")
            self.current_image = self.original_image.copy()
            self.show_image(self.current_image)
        except Exception as e:
            messagebox.showerror("错误", f"载入图像失败：{str(e)}")

    def show_image(self, img):
        if img is None:
            return
        if len(img.shape) == 2:
            img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        else:
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        pil_img = Image.fromarray(img_rgb)
        canvas_w = self.canvas.winfo_width()
        canvas_h = self.canvas.winfo_height()
        img_h, img_w = img.shape[:2]
        scale = min(canvas_w / img_w, canvas_h / img_h)
        show_w, show_h = int(img_w * scale), int(img_h * scale)
        pad_x = (canvas_w - show_w) // 2
        pad_y = (canvas_h - show_h) // 2
        self._img_scale = scale
        self._img_pad_x = pad_x
        self._img_pad_y = pad_y
        pil_img = pil_img.resize((show_w, show_h), Image.LANCZOS)
        tk_img = ImageTk.PhotoImage(pil_img)
        self.canvas.delete("all")
        self.canvas.create_image(pad_x, pad_y, anchor=tk.NW, image=tk_img)
        self.canvas.image = tk_img
        self.tk_img = tk_img

    def ask_iterations(self, title="请输入操作次数", prompt="操作次数：", default=1, minval=1, maxval=20, on_change=None):
        input_win = tk.Toplevel(self.root)
        input_win.title(title)
        input_win.geometry("300x120")
        tk.Label(input_win, text=prompt).pack(pady=5)
        var = tk.IntVar(value=default)
        result = {'value': default}
        def scale_callback(val):
            val = int(val)
            if on_change:
                on_change(val)
        scale = tk.Scale(input_win, from_=minval, to=maxval, orient=tk.HORIZONTAL, variable=var, length=200, command=scale_callback)
        scale.pack(pady=5)
        def on_ok():
            val = var.get()
            if val < minval or val > maxval:
                tk.messagebox.showerror("错误", f"请输入{minval}-{maxval}之间的整数！")
                return
            result['value'] = val
            input_win.destroy()
        tk.Button(input_win, text="确定", command=on_ok).pack(pady=5)
        input_win.transient(self.root)
        input_win.grab_set()
        self.root.wait_window(input_win)
        return result['value']

    def ask_iterations_and_ksize(self, title="闭运算参数选择", prompt1="操作次数：", prompt2="核大小(奇数)：", default_iter=1, min_iter=1, max_iter=20, default_ksize=3, min_ksize=3, max_ksize=15, on_change=None):
        input_win = tk.Toplevel(self.root)
        input_win.title(title)
        input_win.geometry("340x200")
        tk.Label(input_win, text=prompt1).pack(pady=2)
        var_iter = tk.IntVar(value=default_iter)
        var_ksize = tk.IntVar(value=default_ksize)
        result = {'iterations': default_iter, 'ksize': default_ksize}
        def scale_callback(val=None):
            iter_val = var_iter.get()
            ksize_val = var_ksize.get()
            # 保证核大小为奇数
            if ksize_val % 2 == 0:
                ksize_val += 1
                var_ksize.set(ksize_val)
            if on_change:
                on_change(iter_val, ksize_val)
        scale_iter = tk.Scale(input_win, from_=min_iter, to=max_iter, orient=tk.HORIZONTAL, variable=var_iter, length=260, command=lambda v: scale_callback())
        scale_iter.pack(pady=2)
        tk.Label(input_win, text=prompt2).pack(pady=2)
        scale_ksize = tk.Scale(input_win, from_=min_ksize, to=max_ksize, orient=tk.HORIZONTAL, variable=var_ksize, length=260, command=lambda v: scale_callback())
        scale_ksize.pack(pady=2)
        def on_ok():
            iter_val = var_iter.get()
            ksize_val = var_ksize.get()
            if iter_val < min_iter or iter_val > max_iter:
                tk.messagebox.showerror("错误", f"操作次数请输入{min_iter}-{max_iter}之间的整数！")
                return
            if ksize_val % 2 == 0 or ksize_val < min_ksize or ksize_val > max_ksize:
                tk.messagebox.showerror("错误", f"核大小必须为{min_ksize}-{max_ksize}之间的奇数！")
                return
            result['iterations'] = iter_val
            result['ksize'] = ksize_val
            input_win.destroy()
        tk.Button(input_win, text="确定", command=on_ok).pack(pady=5)
        input_win.transient(self.root)
        input_win.grab_set()
        self.root.wait_window(input_win)
        return result['iterations'], result['ksize']

    def process_image(self, operation_id):
        if self.current_image is None:
            tk.messagebox.showwarning("提示", "请先载入图像！")
            return
        processed_img = None
        gray_img = to_gray(self.current_image) if operation_id != 1 else None
        try:
            if operation_id == 1:
                processed_img = to_gray(self.current_image)
            elif operation_id == 2:
                processed_img = otsu_threshold(gray_img)
                processed_img = 255 - processed_img
            elif operation_id == 3:
                # 实时预览腐蚀
                backup = gray_img.copy()
                def preview(val):
                    img = erode(backup, iterations=val)
                    self.show_image(img)
                iterations = self.ask_iterations(title="腐蚀操作次数", prompt="请选择腐蚀操作次数：", default=1, minval=1, maxval=20, on_change=preview)
                processed_img = erode(backup, iterations=iterations)
            elif operation_id == 4:
                # 实时预览膨胀
                backup = gray_img.copy()
                def preview(val):
                    img = dilate(backup, iterations=val)
                    self.show_image(img)
                iterations = self.ask_iterations(title="膨胀操作次数", prompt="请选择膨胀操作次数：", default=1, minval=1, maxval=20, on_change=preview)
                processed_img = dilate(backup, iterations=iterations)
            elif operation_id == 6:
                # 实时预览闭运算
                backup = gray_img.copy()
                def preview(iter_val, ksize_val):
                    bin_img = otsu_threshold(backup)
                    closed = morphology_close(bin_img, ksize=ksize_val, iterations=iter_val)
                    inv = 255 - closed
                    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(inv, connectivity=8)
                    min_area = 30
                    for i in range(1, num_labels):
                        if stats[i, cv2.CC_STAT_AREA] < min_area:
                            inv[labels == i] = 0
                    img = 255 - inv
                    self.show_image(img)
                iterations, ksize = self.ask_iterations_and_ksize(
                    title="闭运算参数选择",
                    prompt1="请选择闭运算操作次数：",
                    prompt2="请选择核大小(奇数)：",
                    default_iter=1, min_iter=1, max_iter=20,
                    default_ksize=3, min_ksize=3, max_ksize=15,
                    on_change=preview
                )
                bin_img = otsu_threshold(backup)
                closed = morphology_close(bin_img, ksize=ksize, iterations=iterations)
                inv = 255 - closed
                num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(inv, connectivity=8)
                min_area = 30
                for i in range(1, num_labels):
                    if stats[i, cv2.CC_STAT_AREA] < min_area:
                        inv[labels == i] = 0
                processed_img = 255 - inv
            elif operation_id == 7:
                processed_img = canny_edge(gray_img)
            else:
                tk.messagebox.showwarning("提示", "未知操作！")
                return
            self.current_image = processed_img
            self.show_image(processed_img)
        except Exception as e:
            tk.messagebox.showerror("错误", f"处理失败：{str(e)}")

    def crop_image(self):
        if self.current_image is None:
            tk.messagebox.showinfo("错误", "没有图像可裁剪")
            return
        if not self.crop_region:
            tk.messagebox.showinfo("错误", "请先用鼠标绘制裁剪区域")
            return
        img = self.current_image
        h, w = img.shape[:2]
        scale = getattr(self, '_img_scale', 1.0)
        pad_x = getattr(self, '_img_pad_x', 0)
        pad_y = getattr(self, '_img_pad_y', 0)
        import math
        y0 = int(math.floor((self.crop_region[0] - pad_y) / scale))
        y1 = int(math.ceil((self.crop_region[1] - pad_y) / scale))
        x0 = int(math.floor((self.crop_region[2] - pad_x) / scale))
        x1 = int(math.ceil((self.crop_region[3] - pad_x) / scale))
        y0 = max(0, min(h-1, y0))
        y1 = max(y0+1, min(h, y1))
        x0 = max(0, min(w-1, x0))
        x1 = max(x0+1, min(w, x1))
        cropped = img[y0:y1, x0:x1]
        self.current_image = cropped
        self.show_image(self.current_image)
        self.crop_region = None
        if self.current_rect:
            self.canvas.delete(self.current_rect)
            self.current_rect = None
        if self.original_image is not None:
            oh, ow = self.original_image.shape[:2]
            y0o = max(0, min(oh-1, y0))
            y1o = max(y0o+1, min(oh, y1))
            x0o = max(0, min(ow-1, x0))
            x1o = max(x0o+1, min(ow, x1))
            self.cropped_original_image = self.original_image[y0o:y1o, x0o:x1o]
        else:
            self.cropped_original_image = None

    def on_mouse_down(self, event):
        if self.current_image is None:
            return
        if self.current_rect and not self.is_drawing:
            self.canvas.delete(self.current_rect)
            self.current_rect = None
        self.rect_start_x = max(0, min(event.x, self.canvas.winfo_width() - 1))
        self.rect_start_y = max(0, min(event.y, self.canvas.winfo_height() - 1))
        self.is_drawing = True

    def on_mouse_move(self, event):
        if self.is_drawing and self.current_image is not None:
            if self.current_rect:
                self.canvas.delete(self.current_rect)
            x1 = max(0, min(event.x, self.canvas.winfo_width() - 1))
            y1 = max(0, min(event.y, self.canvas.winfo_height() - 1))
            self.current_rect = self.canvas.create_rectangle(
                self.rect_start_x, self.rect_start_y, x1, y1, outline="red", width=2)

    def on_mouse_up(self, event):
        if self.is_drawing and self.current_image is not None:
            x1 = max(0, min(event.x, self.canvas.winfo_width() - 1))
            y1 = max(0, min(event.y, self.canvas.winfo_height() - 1))
            self.rect_end_x = x1
            self.rect_end_y = y1
            self.is_drawing = False
            if abs(self.rect_end_x - self.rect_start_x) > 2 and abs(self.rect_end_y - self.rect_start_y) > 2:
                self.crop_region = (
                    min(self.rect_start_y, self.rect_end_y),
                    max(self.rect_start_y, self.rect_end_y),
                    min(self.rect_start_x, self.rect_end_x),
                    max(self.rect_start_x, self.rect_end_x)
                )
            else:
                self.crop_region = None
                if self.current_rect:
                    self.canvas.delete(self.current_rect)
                    self.current_rect = None

    def reset_to_original(self):
        if self.original_image is not None:
            self.current_image = self.original_image.copy()
            self.show_image(self.current_image)
            self.crop_region = None
            self.cropped_original_image = None
        else:
            tk.messagebox.showinfo("提示", "未载入原图，无法重置！")

    def contour_detect_with_canny(self):
        if self.current_image is None:
            tk.messagebox.showwarning("提示", "请先载入图像！")
            return
        if hasattr(self, "cropped_original_image") and self.cropped_original_image is not None:
            base_img = self.cropped_original_image.copy()
        else:
            base_img = self.original_image.copy()
        if len(self.current_image.shape) == 3:
            gray = cv2.cvtColor(self.current_image, cv2.COLOR_BGR2GRAY)
        else:
            gray = self.current_image.copy()
        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        closed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=1)
        closed = cv2.dilate(closed, kernel, iterations=1)
        edges = cv2.Canny(closed, 50, 150)
        contour_img, merged_boxes = find_contours_with_merge(edges, area_thresh=400, wh_ratio=(0.3, 3), min_wh=12, merge_thresh=12)
        self.char_boxes = merged_boxes if merged_boxes else []
        if merged_boxes:
            merged_boxes = sorted(merged_boxes, key=lambda b: b[1]+b[3]//2, reverse=True)
            self.last_char_box = merged_boxes[0]
        else:
            self.last_char_box = None
        if len(base_img.shape) == 2:
            result = cv2.cvtColor(base_img, cv2.COLOR_GRAY2BGR)
        else:
            result = base_img
        for x, y, w, h in merged_boxes:
            cv2.rectangle(result, (x, y), (x + w, y + h), (0, 255, 0), 2)
        self.current_image = result
        self.show_image(result)

    def extract_and_save_chars(self):
        if self.current_image is None:
            messagebox.showwarning("提示", "请先载入图像！")
            return
        if hasattr(self, "cropped_original_image") and self.cropped_original_image is not None:
            base_img = self.cropped_original_image.copy()
        elif self.original_image is not None:
            base_img = self.original_image.copy()
        else:
            base_img = self.current_image
        if not hasattr(self, "char_boxes") or not self.char_boxes:
            messagebox.showinfo("提示", "请先进行轮廓检测，且检测到汉字轮廓！")
            return
        save_dir = os.path.join(os.path.dirname(__file__), 'chars')
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        count = 0
        def sort_boxes(boxes, y_thresh=20):
            boxes = sorted(boxes, key=lambda b: b[1])
            lines = []
            for b in boxes:
                x, y, w, h = b
                found = False
                for line in lines:
                    if abs(line[0][1] - y) < y_thresh:
                        line.append(b)
                        found = True
                        break
                if not found:
                    lines.append([b])
            sorted_boxes = []
            for line in lines:
                line = sorted(line, key=lambda b: b[0])
                sorted_boxes.extend(line)
            return sorted_boxes
        for x, y, w, h in sort_boxes(self.char_boxes):
            if w < 10 or h < 10 or w*h < 80:
                continue
            x0 = max(0, x)
            y0 = max(0, y)
            x1 = min(base_img.shape[1], x + w)
            y1 = min(base_img.shape[0], y + h)
            char_img = base_img[y0:y1, x0:x1]
            if char_img.size == 0:
                continue
            if len(char_img.shape) == 3:
                char_gray = cv2.cvtColor(char_img, cv2.COLOR_BGR2GRAY)
            else:
                char_gray = char_img
            _, char_bin = cv2.threshold(char_gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            size = max(char_bin.shape)
            square = np.ones((size, size), dtype=np.uint8) * 255
            y_offset = (size - char_bin.shape[0]) // 2
            x_offset = (size - char_bin.shape[1]) // 2
            square[y_offset:y_offset+char_bin.shape[0], x_offset:x_offset+char_bin.shape[1]] = char_bin
            char_resized = cv2.resize(square, (64, 64), interpolation=cv2.INTER_AREA)
            save_path = os.path.join(save_dir, f"{count}.png")
            ext = os.path.splitext(save_path)[1]
            success, buf = cv2.imencode(ext, char_resized)
            if success:
                with open(save_path, 'wb') as f:
                    f.write(buf)
                count += 1
        messagebox.showinfo("提示", f"已保存 {count} 个汉字图像到 chars 目录！")

    def apply_median_filter(self):
        if self.current_image is None:
            tk.messagebox.showwarning("提示", "请先载入图像！")
            return
        ksize = self.ask_iterations(title="中值滤波核大小", prompt="请选择奇数核大小(如3,5,7)：", default=3, minval=3, maxval=15)
        if ksize % 2 == 0 or ksize < 1:
            tk.messagebox.showerror("错误", "核大小必须为正奇数！")
            return
        gray_img = to_gray(self.current_image)
        filtered = median_filter(gray_img, ksize=ksize)
        self.current_image = filtered
        self.show_image(filtered)

# ========== 程序入口 ==========
if __name__ == "__main__":
    root = tk.Tk()
    app = ImageProcessingApp(root)
    root.mainloop()
